#!/bin/bash

# Kafka to Elasticsearch Time Partition Job Submission Script
# ===========================================================

# 脚本配置
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
JOB_NAME="kafka2es-timepartition"
JAR_FILE="${PROJECT_ROOT}/target/flink-hbase-1.0-SNAPSHOT.jar"
MAIN_CLASS="com.flink.hbase.kafka2elasticsearch.timepartition.KafkaToElasticsearchTimePartitionJob"

# 默认配置
DEFAULT_FLINK_HOME="/opt/flink"
DEFAULT_CONFIG_FILE="${PROJECT_ROOT}/src/main/resources/kafka2es-timepartition-config.properties"
DEFAULT_PARALLELISM=4
DEFAULT_KAFKA_PARALLELISM=4
DEFAULT_PROCESSING_PARALLELISM=8
DEFAULT_ES_PARALLELISM=2

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_debug() {
    echo -e "${BLUE}[DEBUG]${NC} $1"
}

# 打印使用说明
print_usage() {
    cat << EOF
Usage: $0 [OPTIONS]

Kafka to Elasticsearch Time Partition Job Submission Script

OPTIONS:
    -h, --help                      显示帮助信息
    -c, --config FILE              配置文件路径 (默认: ${DEFAULT_CONFIG_FILE})
    -f, --flink-home PATH          Flink 安装路径 (默认: ${DEFAULT_FLINK_HOME})
    -j, --jar FILE                 JAR 文件路径 (默认: ${JAR_FILE})
    
    # Kafka 配置
    --kafka-brokers BROKERS        Kafka bootstrap servers (如: localhost:9092)
    --kafka-topics TOPICS          Kafka topics，逗号分隔 (如: topic1,topic2)
    --kafka-group-id GROUP_ID      Kafka consumer group ID
    --kafka-offset-reset RESET     Kafka offset reset策略 (earliest|latest|committed)
    
    # Elasticsearch 配置
    --es-hosts HOSTS               ES 主机列表 (如: localhost:9200)
    --es-username USERNAME         ES 用户名
    --es-password PASSWORD         ES 密码
    --index-name INDEX_NAME        基础索引名称
    --document-id-field FIELD      文档 ID 字段名
    
    # 时间处理配置
    --time-fields FIELDS           时间字段列表，逗号分隔 (如: timestamp,create_time)
    --timezone TIMEZONE           时区设置 (如: Asia/Shanghai)
    
    # 并行度配置
    --parallelism NUM              全局并行度 (默认: ${DEFAULT_PARALLELISM})
    --kafka-parallelism NUM        Kafka 消费者并行度 (默认: ${DEFAULT_KAFKA_PARALLELISM})
    --processing-parallelism NUM   处理并行度 (默认: ${DEFAULT_PROCESSING_PARALLELISM})
    --es-parallelism NUM           ES 写入并行度 (默认: ${DEFAULT_ES_PARALLELISM})
    
    # 批处理配置
    --es-batch-size SIZE           ES 批量写入大小 (默认: 500)
    --es-flush-interval MS         ES 刷新间隔，毫秒 (默认: 5000)
    
    # Checkpoint 配置
    --checkpoint-interval MS       Checkpoint 间隔，毫秒 (默认: 60000)
    
    # 其他选项
    --dry-run                      仅显示命令，不执行
    --verbose                      详细输出
    --optimize                     显示优化建议

EXAMPLES:
    # 使用默认配置文件
    $0 -c kafka2es-config.properties
    
    # 使用命令行参数
    $0 --kafka-brokers localhost:9092 \\
       --kafka-topics user_events \\
       --kafka-group-id es-consumer \\
       --es-hosts localhost:9200 \\
       --index-name user_events \\
       --kafka-parallelism 8 \\
       --es-parallelism 2
    
    # 获取优化建议
    $0 --optimize
    
    # 测试运行
    $0 --dry-run --verbose

EOF
}

# 检查环境
check_environment() {
    log_info "检查环境配置..."
    
    # 检查 Flink 安装
    if [[ ! -d "$FLINK_HOME" ]]; then
        log_error "Flink 安装目录不存在: $FLINK_HOME"
        log_info "请设置 FLINK_HOME 环境变量或使用 --flink-home 参数"
        exit 1
    fi
    
    # 检查 flink 命令
    if [[ ! -x "$FLINK_HOME/bin/flink" ]]; then
        log_error "Flink 命令不存在: $FLINK_HOME/bin/flink"
        exit 1
    fi
    
    # 检查 JAR 文件
    if [[ ! -f "$JAR_FILE" ]]; then
        log_error "JAR 文件不存在: $JAR_FILE"
        log_info "请先构建项目: mvn clean package"
        exit 1
    fi
    
    log_info "环境检查通过"
}

# 检查 Kafka 连接
check_kafka_connection() {
    local brokers="$1"
    if [[ -z "$brokers" ]]; then
        log_warn "未指定 Kafka brokers，跳过连接检查"
        return 0
    fi
    
    log_info "检查 Kafka 连接: $brokers"
    
    # 这里可以添加 Kafka 连接检查逻辑
    # 例如使用 kafka-topics.sh 命令
    
    log_info "Kafka 连接检查通过"
}

# 检查 Elasticsearch 连接
check_elasticsearch_connection() {
    local hosts="$1"
    if [[ -z "$hosts" ]]; then
        log_warn "未指定 ES hosts，跳过连接检查"
        return 0
    fi
    
    log_info "检查 Elasticsearch 连接: $hosts"
    
    # 简单的 HTTP 检查
    if command -v curl &> /dev/null; then
        local es_host=$(echo "$hosts" | cut -d',' -f1)
        if curl -s "http://$es_host/_cluster/health" &> /dev/null; then
            log_info "Elasticsearch 连接检查通过"
        else
            log_warn "Elasticsearch 连接检查失败，但作业仍将提交"
        fi
    else
        log_warn "未安装 curl，跳过 ES 连接检查"
    fi
}

# 获取 Topic 分区信息
get_topic_partitions() {
    local brokers="$1"
    local topics="$2"
    
    if [[ -z "$brokers" || -z "$topics" ]]; then
        return 0
    fi
    
    log_info "获取 Topic 分区信息..."
    
    # 这里可以添加获取分区信息的逻辑
    # 例如使用 kafka-topics.sh --describe
    
    log_info "建议根据 Topic 分区数设置 Kafka 并行度"
}

# 显示优化建议
show_optimization_recommendations() {
    cat << EOF

${GREEN}=== Kafka 到 Elasticsearch 时间分区作业优化建议 ===${NC}

${BLUE}1. Kafka 并行度优化:${NC}
   - Kafka 消费者并行度应该等于 Topic 分区数
   - 获取分区信息: kafka-topics --bootstrap-server localhost:9092 --describe --topic <topic-name>
   - 建议: --kafka-parallelism <partition-count>

${BLUE}2. 处理并行度优化:${NC}
   - 处理并行度建议为 CPU 核心数的 2-4 倍
   - 建议: --processing-parallelism <cpu-cores * 2-4>

${BLUE}3. Elasticsearch 并行度优化:${NC}
   - ES 并行度建议为 ES 节点数或稍少
   - 建议: --es-parallelism <es-node-count>

${BLUE}4. 批处理优化:${NC}
   - 批量大小: 100-1000 条记录
   - 刷新间隔: 5-30 秒
   - 建议: --es-batch-size 500 --es-flush-interval 10000

${BLUE}5. 时间分区优化:${NC}
   - 优先使用时间戳字段而非字符串解析
   - 设置正确的时区: --timezone Asia/Shanghai
   - 时间字段优先级: timestamp > create_time > event_time

${BLUE}6. 稳定性配置:${NC}
   - 启用 Checkpoint: --checkpoint-interval 60000
   - 使用 RocksDB 状态后端处理大状态
   - 配置合适的重启策略

${BLUE}7. 资源配置建议:${NC}
   - JVM 堆内存: 物理内存的 70-80%
   - 网络缓冲区: 128MB - 1GB
   - 任务槽数: CPU 核心数

${BLUE}8. 监控指标:${NC}
   - Kafka 消费延迟 (Consumer Lag)
   - ES 索引速率和错误率
   - Flink 背压和 Checkpoint 时间

EOF
}

# 构建 Flink 提交命令
build_flink_command() {
    local cmd="$FLINK_HOME/bin/flink run"
    
    # 添加 JAR 文件
    cmd="$cmd $JAR_FILE"
    
    # 添加参数
    if [[ -n "$KAFKA_BROKERS" ]]; then
        cmd="$cmd --kafka-brokers $KAFKA_BROKERS"
    fi
    
    if [[ -n "$KAFKA_TOPICS" ]]; then
        cmd="$cmd --kafka-topics $KAFKA_TOPICS"
    fi
    
    if [[ -n "$KAFKA_GROUP_ID" ]]; then
        cmd="$cmd --kafka-group-id $KAFKA_GROUP_ID"
    fi
    
    if [[ -n "$KAFKA_OFFSET_RESET" ]]; then
        cmd="$cmd --kafka-offset-reset $KAFKA_OFFSET_RESET"
    fi
    
    if [[ -n "$ES_HOSTS" ]]; then
        cmd="$cmd --es-hosts $ES_HOSTS"
    fi
    
    if [[ -n "$ES_USERNAME" ]]; then
        cmd="$cmd --es-username $ES_USERNAME"
    fi
    
    if [[ -n "$ES_PASSWORD" ]]; then
        cmd="$cmd --es-password $ES_PASSWORD"
    fi
    
    if [[ -n "$INDEX_NAME" ]]; then
        cmd="$cmd --index-name $INDEX_NAME"
    fi
    
    if [[ -n "$DOCUMENT_ID_FIELD" ]]; then
        cmd="$cmd --document-id-field $DOCUMENT_ID_FIELD"
    fi
    
    if [[ -n "$TIME_FIELDS" ]]; then
        cmd="$cmd --time-fields $TIME_FIELDS"
    fi
    
    if [[ -n "$TIMEZONE" ]]; then
        cmd="$cmd --timezone $TIMEZONE"
    fi
    
    if [[ -n "$KAFKA_PARALLELISM" ]]; then
        cmd="$cmd --kafka-parallelism $KAFKA_PARALLELISM"
    fi
    
    if [[ -n "$PROCESSING_PARALLELISM" ]]; then
        cmd="$cmd --processing-parallelism $PROCESSING_PARALLELISM"
    fi
    
    if [[ -n "$ES_PARALLELISM" ]]; then
        cmd="$cmd --es-parallelism $ES_PARALLELISM"
    fi
    
    if [[ -n "$ES_BATCH_SIZE" ]]; then
        cmd="$cmd --es-batch-size $ES_BATCH_SIZE"
    fi
    
    if [[ -n "$ES_FLUSH_INTERVAL" ]]; then
        cmd="$cmd --es-flush-interval $ES_FLUSH_INTERVAL"
    fi
    
    if [[ -n "$CHECKPOINT_INTERVAL" ]]; then
        cmd="$cmd --checkpoint-interval $CHECKPOINT_INTERVAL"
    fi
    
    if [[ -n "$CONFIG_FILE" ]]; then
        cmd="$cmd --config-file $CONFIG_FILE"
    fi
    
    echo "$cmd"
}

# 主函数
main() {
    # 初始化变量
    FLINK_HOME="${DEFAULT_FLINK_HOME}"
    CONFIG_FILE="${DEFAULT_CONFIG_FILE}"
    DRY_RUN=false
    VERBOSE=false
    OPTIMIZE=false
    
    # 解析命令行参数
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                print_usage
                exit 0
                ;;
            -c|--config)
                CONFIG_FILE="$2"
                shift 2
                ;;
            -f|--flink-home)
                FLINK_HOME="$2"
                shift 2
                ;;
            -j|--jar)
                JAR_FILE="$2"
                shift 2
                ;;
            --kafka-brokers)
                KAFKA_BROKERS="$2"
                shift 2
                ;;
            --kafka-topics)
                KAFKA_TOPICS="$2"
                shift 2
                ;;
            --kafka-group-id)
                KAFKA_GROUP_ID="$2"
                shift 2
                ;;
            --kafka-offset-reset)
                KAFKA_OFFSET_RESET="$2"
                shift 2
                ;;
            --es-hosts)
                ES_HOSTS="$2"
                shift 2
                ;;
            --es-username)
                ES_USERNAME="$2"
                shift 2
                ;;
            --es-password)
                ES_PASSWORD="$2"
                shift 2
                ;;
            --index-name)
                INDEX_NAME="$2"
                shift 2
                ;;
            --document-id-field)
                DOCUMENT_ID_FIELD="$2"
                shift 2
                ;;
            --time-fields)
                TIME_FIELDS="$2"
                shift 2
                ;;
            --timezone)
                TIMEZONE="$2"
                shift 2
                ;;
            --parallelism)
                PARALLELISM="$2"
                shift 2
                ;;
            --kafka-parallelism)
                KAFKA_PARALLELISM="$2"
                shift 2
                ;;
            --processing-parallelism)
                PROCESSING_PARALLELISM="$2"
                shift 2
                ;;
            --es-parallelism)
                ES_PARALLELISM="$2"
                shift 2
                ;;
            --es-batch-size)
                ES_BATCH_SIZE="$2"
                shift 2
                ;;
            --es-flush-interval)
                ES_FLUSH_INTERVAL="$2"
                shift 2
                ;;
            --checkpoint-interval)
                CHECKPOINT_INTERVAL="$2"
                shift 2
                ;;
            --dry-run)
                DRY_RUN=true
                shift
                ;;
            --verbose)
                VERBOSE=true
                shift
                ;;
            --optimize)
                OPTIMIZE=true
                shift
                ;;
            *)
                log_error "未知参数: $1"
                print_usage
                exit 1
                ;;
        esac
    done
    
    # 显示优化建议
    if [[ "$OPTIMIZE" == true ]]; then
        show_optimization_recommendations
        exit 0
    fi
    
    # 检查环境
    check_environment
    
    # 检查连接
    check_kafka_connection "$KAFKA_BROKERS"
    check_elasticsearch_connection "$ES_HOSTS"
    
    # 获取 Topic 分区信息
    get_topic_partitions "$KAFKA_BROKERS" "$KAFKA_TOPICS"
    
    # 构建命令
    FLINK_CMD=$(build_flink_command)
    
    # 显示信息
    log_info "准备提交 Kafka 到 Elasticsearch 时间分区作业"
    log_info "JAR 文件: $JAR_FILE"
    log_info "主类: $MAIN_CLASS"
    
    if [[ "$VERBOSE" == true ]]; then
        log_debug "完整命令: $FLINK_CMD"
    fi
    
    # 执行或显示命令
    if [[ "$DRY_RUN" == true ]]; then
        log_info "DRY RUN - 以下命令将被执行:"
        echo "$FLINK_CMD"
    else
        log_info "提交作业..."
        eval "$FLINK_CMD"
        
        if [[ $? -eq 0 ]]; then
            log_info "作业提交成功"
            log_info "可以使用 Flink Web UI 监控作业状态: http://localhost:8081"
        else
            log_error "作业提交失败"
            exit 1
        fi
    fi
}

# 执行主函数
main "$@" 